import cn.hutool.core.io.resource.ClassPathResource;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import service.FieldMappingService;
import service.FieldMappingServiceDefaultImpl;
import util.CommonUtil;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.Statement;
import java.util.HashMap;
import java.util.Map;
import java.util.Scanner;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 从数据库导入数据到elastic Search
 */
@Slf4j(topic = "m.DataImport2ElasticSearch")
public class DataImport2ElasticSearch {

    /**
     * 数据库链接
     */
    private static Connection connection;
    /**
     * 操作es客户端
     */
    private static RestHighLevelClient restHighLevelClient;

    public static void main(String[] args) throws IOException {
        try {
            final Scanner scanner = new Scanner(System.in);
            log.info("***********数据库导入start*************");
            log.info("*************初始化database-start*************");
            final AtomicInteger fetchSize = new AtomicInteger(1000);
            connection = CommonUtil.initDataBase(log, scanner, fetchSize);
            log.info("*************初始化database-成功*************");
            log.info("*************初始化elasticSearch-start*************");
            restHighLevelClient = CommonUtil.initElasticSearchClient(log, scanner);
            log.info("*************初始化elasticSearch-成功*************");

            //获取索引index
            log.info("请输入要写入的索引名称（请先确定elastic search中有索引，并且mapping已设置！！！）");
            String tmp = CommonUtil.receiveLine(log, scanner);
            assertNull(tmp);
            final String esIndex = tmp;
            IndexRequest indexRequest = new IndexRequest(esIndex);

            //创建statement
            Statement statement = connection.createStatement();
            //设置fetch_size大小
            statement.setFetchSize(fetchSize.get());
            //查询数据库
            log.info("请输入查询SQL文件路径（注意：文件中不要有多条sql、不要【;】）");
            log.info("如若使用包内文件，前面加classpath:");
            tmp = CommonUtil.receiveLine(log, scanner);
            assertNull(tmp);
            File sqlFile = null;
            String sql = null;
            if (tmp.startsWith("classpath")) {
                String[] split = tmp.split(":");
                ClassPathResource classPathResource = new ClassPathResource(split[1]);
                BufferedReader reader = classPathResource.getReader(StandardCharsets.UTF_8);
                sql = reader.lines().collect(Collectors.joining("\n"));

            } else {
                sqlFile = new File(tmp);
                if (!sqlFile.exists()) {
                    throw new Exception("sql文件不存在");
                }
                cn.hutool.core.io.file.FileReader fileReader = new cn.hutool.core.io.file.FileReader(sqlFile);
                sql = fileReader.readString();
            }

            final String countSql = "select count(1) from (" + sql + ") a";
            log.info("正在执行的SQL：" + countSql);
            ResultSet resultSet = statement.executeQuery(countSql);
            final int count = resultSet.next() ? resultSet.getInt(1) : 0;
            log.info("查询总数:" + count);
            log.info("正在执行的SQL：" + sql);
            resultSet = statement.executeQuery(sql);
            //获取返回结果集相关信息
            final ResultSetMetaData metaData = resultSet.getMetaData();
            //获取表一共几列？
            final int columnCount = metaData.getColumnCount();
            //生命特殊字段处理类
            log.info("请输入自定义处理表字段映射关系的业务类（没有可以不填,必须实现接口FieldMappingService）");
            tmp = CommonUtil.receiveLine(log, scanner);
            FieldMappingService fieldMappingService = new FieldMappingServiceDefaultImpl();
            if (tmp != null && !tmp.equals("")) {
                final Class<?> aClass = Class.forName(tmp);
                fieldMappingService = (FieldMappingService) aClass.newInstance();
            }


            //设置字段值对应关系
            final Map<String, String> fieldMapping = new HashMap<>();
            int dataInt = 1;
            while (resultSet.next()) {
                String result = null;
                try {
                    final Map<String, Object> row = new HashMap<>();
                    for (int i = 1; i <= columnCount; i++) {
                        final Object object = resultSet.getObject(i);
                        final String columnName = metaData.getColumnName(i);
                        if (!fieldMapping.containsKey(columnName)) {
                            log.info("请输入字段对应关系，当前字段为：" + columnName + "，需要对应es字段为：(不填即默认)");
                            tmp = CommonUtil.receiveLine(log, scanner);
                            if (StringUtils.isBlank(tmp)) {
                                tmp = columnName;
                            }
                            fieldMapping.put(columnName, tmp);
                            log.info("字段对应关系为：" + JSON.toJSONString(fieldMapping));
                            row.put(fieldMapping.get(columnName), fieldMappingService.getValueByField(columnName, object));

                            if (!fieldMapping.containsKey("elasticUniId")) {
                                log.info("此字段是elasticSearch的主键id吗？0:不是|1:是（注：如果最终不设置主键id，那es会自己生成随机id）");
                                tmp = CommonUtil.receiveLine(log, scanner);
                                if (StringUtils.equals(tmp, "1")) {
                                    fieldMapping.put("elasticUniId", fieldMapping.get(columnName));
                                }
                            }
                        } else {
                            row.put(fieldMapping.get(columnName), fieldMappingService.getValueByField(columnName, object));
                        }
                    }
                    if (fieldMapping.containsKey("elasticUniId")) {
                        indexRequest.id(String.valueOf(row.get(fieldMapping.get("elasticUniId"))));
                    }
                    indexRequest.source(row);
                    IndexResponse indexResponse = restHighLevelClient.index(indexRequest, RequestOptions.DEFAULT);
                    result = indexResponse.getId() + "-->" + (indexResponse.getResult() != null ? indexResponse.getResult().toString() : "");
                } catch (Exception e) {
                    e.printStackTrace();
                    log.error("1--{}", ExceptionUtils.getMessage(e));
                } finally {
                    log.info("当前查询总数为" + count + "条，现在正在处理第" + dataInt + "条,结果为：" + result);
                }

                dataInt++;
            }
        } catch (Exception e) {
            e.printStackTrace();
            log.error("2--{}", ExceptionUtils.getMessage(e));
        }
        System.out.println("*******完成*********");
    }

    /**
     * 校验输入空值
     */
    public static void assertNull(String tmp) throws Exception {
        CommonUtil.assertNull(tmp);
    }


}
